#!/usr/bin/env python3
"""
Preprocess the DTU MVS dataset.

This script processes images from the DTU dataset, converting camera parameters,
resizing images and depth maps, and saving the processed outputs into a
directory structured to match the Hypersim processed format.

Usage:
    python preprocess_dtu.py --dtu_dir /path/to/dtu_training --output_dir /path/to/processed_dtu
"""
import os
import os.path as osp
import re
import shutil
import sys
from tqdm import tqdm
import numpy as np
import argparse

os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2

# Assuming these utils are in a reachable path
from utils.parallel import parallel_threads
from utils.cropping import crop_resize_if_necessary


def get_parser():
    """Sets up the argument parser."""
    parser = argparse.ArgumentParser(
        description="Preprocess the DTU dataset by converting camera parameters, "
        "resizing images, and saving processed outputs in Hypersim format."
    )
    parser.add_argument(
        "--dtu_dir",
        required=True,
        help="Root directory of the DTU dataset (e.g., 'dtu_training').",
    )
    parser.add_argument(
        "--output_dir",
        default="data/dtu_processed",
        help="Output directory for processed DTU data.",
    )
    return parser


def main(dtu_root, output_dir):
    """Main processing function."""
    print(">> Listing all sequences from the DTU dataset")
    rectified_dir = osp.join(dtu_root, "Rectified")
    sequences = [f for f in os.listdir(rectified_dir) if osp.isdir(osp.join(rectified_dir, f))]
    
    assert sequences, f"Did not find any sequences at {rectified_dir}"
    print(f"   (found {len(sequences)} sequences)")

    for seq in tqdm(sequences, desc="Processing Scenes"):
        scene_outdir = osp.join(output_dir, seq)
        os.makedirs(scene_outdir, exist_ok=True)

        img_dir = osp.join(rectified_dir, seq)
        img_files = [f for f in os.listdir(img_dir) if f.endswith(".png")]
        
        for img_file in tqdm(img_files, leave=False, desc=f"Processing {seq}"):
            process_single_view(dtu_root, seq, img_file, scene_outdir)

    print(f">> Done, saved everything in {output_dir}/")


def process_single_view(dtu_root, scene_id, img_filename, out_dir):
    """Processes a single view from a DTU scene."""
    img_name_base = img_filename.replace(".png", "")
    
    out_cam_path = osp.join(out_dir, f"{img_name_base}_cam.npz")
    if osp.isfile(out_cam_path):
        return  # Already processed

    # Extract view index from filename
    view_match = re.search(r"rect_(\d+)_", img_filename)
    if not view_match:
        return
    view_idx = int(view_match.group(1)) - 1  # DTU is 1-indexed

    # Construct paths
    rgb_path = osp.join(dtu_root, "Rectified", scene_id, img_filename)
    depth_path = osp.join(dtu_root, "Depths", scene_id, f"depth_map_{view_idx:04d}.pfm")
    cam_path = osp.join(dtu_root, "Cameras_train", f"{view_idx:08d}_cam.txt")

    if not all(osp.exists(p) for p in [rgb_path, depth_path, cam_path]):
        return

    # Load data
    try:
        intrinsics, T_cam2world = _load_pose(cam_path, ret_44=True)
        
        color_image = cv2.cvtColor(
            cv2.imread(rgb_path, cv2.IMREAD_COLOR),
            cv2.COLOR_BGR2RGB,
        )
        color_image = cv2.cvtColor(cv2.imread(rgb_path, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
        depthmap = load_pfm_file(depth_path)
        depthmap[~np.isfinite(depthmap)] = 0
    except Exception as e:
        print(f"Error loading data for {rgb_path}: {e}")
        return

    # Resize image and update intrinsics
    try:
        # The crop_resize_if_necessary function expects a PIL image
        image, depth, intrinsics_out = crop_resize_if_necessary(
            color_image, depthmap, intrinsics, (160, 128) # Note: (H, W)
        )
        print(image.shape)
    except Exception as e:
        print(f"Error processing {img_filename}: {e}")
        return

    # Save in Hypersim format
    out_rgb_path = osp.join(out_dir, f"{img_name_base}_rgb.png")
    out_depth_path = osp.join(out_dir, f"{img_name_base}_depth.exr")
    
    image.save(out_rgb_path)
    cv2.imwrite(out_depth_path, depth.astype(np.float32))
    
    # Save camera parameters
    np.savez(
        out_cam_path,
        intrinsics=intrinsics_out.astype(np.float32),
        pose=T_cam2world.astype(np.float32),
    )


def _load_pose(path, ret_44=False):
    f = open(path)
    RT = np.loadtxt(f, skiprows=1, max_rows=4, dtype=np.float32)
    assert RT.shape == (4, 4)
    RT = np.linalg.inv(RT)  # world2cam to cam2world

    K = np.loadtxt(f, skiprows=2, max_rows=3, dtype=np.float32)
    assert K.shape == (3, 3)

    if ret_44:
        return K, RT
    return K, RT[:3, :3], RT[:3, 3]


def load_pfm_file(file_path):
    with open(file_path, "rb") as file:
        header = file.readline().decode("UTF-8").strip()

        if header == "PF":
            is_color = True
        elif header == "Pf":
            is_color = False
        else:
            raise ValueError("The provided file is not a valid PFM file.")

        # dimensions = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("UTF-8"))
        dimensions_line = file.readline().decode("UTF-8").strip()
        dimensions = re.match(r"^(\d+)\s+(\d+)$", dimensions_line)

        if dimensions:
            img_width, img_height = map(int, dimensions.groups())
        else:
            raise ValueError("Invalid PFM header format.")

        endian_scale = float(file.readline().decode("UTF-8").strip())
        if endian_scale < 0:
            dtype = "<f"  # little-endian
        else:
            dtype = ">f"  # big-endian

        try:
            data_buffer = file.read()
            img_data = np.frombuffer(data_buffer, dtype=dtype)
        except Exception as e:
            print(f"Error processing {file_path}: {e}")
            sys.exit(0)


        if is_color:
            img_data = np.reshape(img_data, (img_height, img_width, 3))
        else:
            img_data = np.reshape(img_data, (img_height, img_width))

        img_data = cv2.flip(img_data, 0)

    return img_data


if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    main(args.dtu_dir, args.output_dir)
